import os

from torch.utils.data import Dataset
from PIL import Image
import numpy as np
import torch
import matplotlib.pyplot as plt

# For test
from torchvision import transforms
from torch.utils.data import DataLoader

# MEAN_A = np.array([113.40, 114.08, 116.45])
# STD_A = np.array([48.30,  46.27,  48.14])
# MEAN_B = np.array([111.07, 114.04, 118.18])
# STD_B = np.array([49.41,  47.01,  47.94])

MEAN_A = np.array([0.0, 0.0, 0.0])
STD_A = np.array([1.0, 1.0, 1.0])
MEAN_B = np.array([0.0, 0.0, 0.0])
STD_B = np.array([1.0, 1.0, 1.0])

num_classes = 7
COLORMAP = [[255,255,255], [128,128,128], [0,128,0], [0,255,0], [128,0,0], [255,0,0], [0,0,255]]
CLASSES = ['unchanged', 'ground', 'low vegetation', 'tree', 'building', 'sports field', 'water']

def colormap2label():
    """构建从RGB到SECOND类别索引的映射"""
    colormap2label = np.zeros(256 ** 3, dtype=np.compat.long)
    for i, colormap in enumerate(COLORMAP):
        colormap2label[
            (colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
    return colormap2label

def lbl2indx(colormap, colormap2label):
    """将SECOND标签中的RGB值映射到它们的类别索引"""
    colormap = colormap.astype('int32')
    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256
           + colormap[:, :, 2])
    return colormap2label[idx]



class BaseDataset(Dataset):
    def __init__(self, base_dir=None, split='train', num=None, sub_set=None, transform=None):
        self._base_dir = base_dir
        self.split = split
        self.transform = transform
        self.sub_set = sub_set
        self.num = num
        if self.split == 'train':
            self.sample_list = os.listdir(os.path.join(self._base_dir, sub_set, 'im1'))
            self.split_dir = os.path.join(self._base_dir, sub_set)
            if self.num is not None:
                self.sample_list = self.sample_list[0:self.num]
        elif self.split == 'test':
            self.sample_list = os.listdir(os.path.join(self._base_dir, 'test', sub_set, 'im1'))  # Is Same
            self.split_dir = os.path.join(self._base_dir, 'test', sub_set)
        print("total {} sample".format(len(self.sample_list)))

    def __len__(self):
        return len(self.sample_list)

    def __getitem__(self, idx):
        case = self.sample_list[idx]
        img_A = Image.open(os.path.join(self.split_dir, 'im1', case))
        img_B = Image.open(os.path.join(self.split_dir, 'im2', case))
        label_A = Image.open(os.path.join(self.split_dir, 'label1', case))
        label_B = Image.open(os.path.join(self.split_dir, 'label2', case))

        # Change the label of different sub set. label.shape: [512, 512]
        label_A, label_B = np.array(label_A, dtype=np.uint8), np.array(label_B, dtype=np.uint8)
        label_A, label_B = lbl2indx(label_A, colormap2label()), lbl2indx(label_B, colormap2label())


        if (self.sub_set == 'sub_set1'):
            label_A[label_A > 2] = 0
            label_B[label_B > 2] = 0

        if (self.sub_set == 'sub_set2'):
            label_A[np.logical_and(label_A > 0, label_A < 3)] = 0
            label_A[np.logical_and(label_A > 4, label_A < 7)] = 0

            label_B[np.logical_and(label_B > 0, label_B < 3)] = 0
            label_B[np.logical_and(label_B > 4, label_B < 7)] = 0

        if (self.sub_set == 'sub_set3'):
            label_A[label_A < 5] = 0
            label_B[label_B < 5] = 0

        # print('label_A.shape:', label_A.shape, 'label_A.unique():', np.unique(label_A))
        # print('label_B.shape:', label_B.shape, 'label_B.unique():', np.unique(label_B))

        sample = {'image_A': img_A, 'label_A': label_A, 'image_B': img_B, 'label_B': label_B}
        # if self.split == "train":
        sample = self.transform(sample)
        sample["idx"] = idx
        return sample


class RandomCrop(object):
    def __init__(self, output_size):
        self.output_size = output_size

    def __call__(self, sample):
        img_A, img_B, lbl_A, lbl_B = sample['image_A'], sample['image_B'], sample['label_A'], sample['label_B']
        img_A = np.array(img_A, np.float32)
        img_A = (img_A - MEAN_A) / STD_A
        img_A = img_A.transpose(2, 0, 1)

        img_B = np.array(img_B, np.float32)
        img_B = (img_B - MEAN_B) / STD_B
        img_B = img_B.transpose(2, 0, 1)


        _, h, w = img_A.shape   # 512, 512
        x_l = 0
        y_l = 0
        if (w - self.output_size[0] > 0):
            x_l = np.random.randint(0, w - self.output_size[0])
            y_l = np.random.randint(0, h - self.output_size[1])

        x_r = x_l + self.output_size[0]
        y_r = y_l + self.output_size[1]

        img_A, img_B = img_A[:, y_l:y_r, x_l:x_r], img_B[:, y_l:y_r, x_l:x_r]
        lbl_A, lbl_B = lbl_A[y_l:y_r, x_l:x_r], lbl_B[y_l:y_r, x_l:x_r]

        img_A, img_B = torch.from_numpy(img_A.astype(np.float32)), torch.from_numpy(img_B.astype(np.float32))
        lbl_A, lbl_B = torch.from_numpy(lbl_A.astype(np.long)), torch.from_numpy(lbl_B.astype(np.long))

        sample = {'image_A': img_A, 'image_B': img_B, 'label_A': lbl_A, 'label_B': lbl_B}
        return sample


'''Sticker dataset loader and transform
'''
class StickerDataset(Dataset):
    def __init__(self, base_dir=None, split='train', num=None, sub_set=None, transform=None):
        self._base_dir = base_dir
        self.split = split
        self.transform = transform
        self.sub_set = sub_set
        self.num = num
        if self.split == 'train':
            self.sample_list = os.listdir(os.path.join(self._base_dir, sub_set, 'im1'))
            self.split_dir = os.path.join(self._base_dir, sub_set)
            if self.num is not None:
                self.sample_list = self.sample_list[0:self.num]
        elif self.split == 'test':
            self.sample_list = os.listdir(os.path.join(self._base_dir, 'test', sub_set, 'im1'))  # Is Same
            self.split_dir = os.path.join(self._base_dir, 'test', sub_set)
        print("total {} sample".format(len(self.sample_list)))

    def __len__(self):
        return len(self.sample_list)

    def __getitem__(self, idx):
        case = self.sample_list[idx]
        img_A = Image.open(os.path.join(self.split_dir, 'im1', case))
        img_B = Image.open(os.path.join(self.split_dir, 'im2', case))
        label = Image.open(os.path.join(self.split_dir, 'label', case))


        # Change the label of different sub set. label.shape: [512, 512]
        label = np.array(label, dtype=np.uint8)
        label[label > 0] = 1
        # print('label.shape:', label.shape, 'label.unique():', np.unique(label))

        sample = {'image_A': img_A, 'image_B': img_B, 'label': label}
        # if self.split == "train":
        sample = self.transform(sample)
        sample["idx"] = idx
        return sample


class StickerRandomCrop(object):
    def __init__(self, output_size):
        self.output_size = output_size

    def __call__(self, sample):
        img_A, img_B, lbl = sample['image_A'], sample['image_B'], sample['label']
        img_A = np.array(img_A, np.float32)
        # img_A = np.array(img_A)
        # img_A = (img_A - MEAN_A) / STD_A
        img_A = img_A.transpose(2, 0, 1)


        img_B = np.array(img_B, np.float32)
        # img_B = (img_B - MEAN_B) / STD_B
        img_B = img_B.transpose(2, 0, 1)


        _, h, w = img_A.shape   # 512, 512
        x_l = 0
        y_l = 0
        if (w - self.output_size[0] > 0):
            x_l = np.random.randint(0, w - self.output_size[0])
            y_l = np.random.randint(0, h - self.output_size[1])

        x_r = x_l + self.output_size[0]
        y_r = y_l + self.output_size[1]

        img_A, img_B = img_A[:, y_l:y_r, x_l:x_r], img_B[:, y_l:y_r, x_l:x_r]
        lbl = lbl[y_l:y_r, x_l:x_r]

        img_A, img_B = torch.from_numpy(img_A.astype(np.float32)), torch.from_numpy(img_B.astype(np.float32))
        lbl = torch.from_numpy(lbl.astype(np.long))

        sample = {'image_A': img_A, 'image_B': img_B, 'label': lbl}
        return sample





if __name__ == '__main__':
    train_set = BaseDataset(base_dir='../../data/SECOND', split='train', sub_set='sub_set1', transform=transforms.Compose([
        RandomCrop([512, 512])
    ]))
    trainloader = DataLoader(train_set, batch_size=4, shuffle=True,
                             num_workers=0, pin_memory=True)

    for i_batch, sampled_batch in enumerate(trainloader):
        print(i_batch)

